package org.keyczar.rsa;

import org.keyczar.exceptions.KeyczarException;
import org.keyczar.interfaces.EncryptingStream;
import org.keyczar.interfaces.SigningStream;
import org.keyczar.interfaces.Stream;

import java.nio.ByteBuffer;
import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import java.security.Signature;
import java.security.interfaces.RSAPublicKey;

import javax.crypto.Cipher;
import javax.crypto.ShortBufferException;

/**
 * {@link Stream} implementation used by {@link RsaPublicKey}.
 */
class RsaEncryptingStream extends RsaVerifyingStream implements EncryptingStream {
  private final Cipher cipher;
  RsaEncryptingStream(Signature signature, Cipher cipher, RSAPublicKey jcePublicKey) {
    super(signature, jcePublicKey);
    this.cipher = cipher;
  }
  
  @Override
  public int doFinalEncrypt(ByteBuffer input, ByteBuffer output)
      throws KeyczarException {
    try {
      final int ciphertextSize = cipher.getOutputSize(input.limit());
      final int outputCapacity = output.limit() - output.position();

      ByteBuffer tmpOutput = ByteBuffer.allocate(ciphertextSize);
      cipher.doFinal(input, tmpOutput);

      if (ciphertextSize == outputCapacity) {
        output.put(tmpOutput.array());

      } else if (ciphertextSize == (outputCapacity + 1)
          && tmpOutput.array()[ciphertextSize - 1] == 0) {
        // There exists at least one JCE (the one IBM ships with some versions of
        // Websphere) which outputs ciphertext that's one byte too long, appending
        // a trailing zero.  We need to trim this byte.
        output.put(tmpOutput.array(), 0, outputCapacity);

      } else {
        throw new KeyczarException("Expected " + outputCapacity + " bytes from encryption "
            + "operation but got " + ciphertextSize);
      }

      return outputCapacity;
    } catch (GeneralSecurityException e) {
      throw new KeyczarException(e);
    }
  }

  @Override
  public SigningStream getSigningStream() {
    return new SigningStream() {
      @Override
      public int digestSize() {
        return 0;
      }

      @Override
      public void initSign() {
        // Do nothing
      }

      @Override
      public void sign(ByteBuffer output) {
        // Do nothing
      }

      @Override
      public void updateSign(ByteBuffer input) {
        // Do nothing
      }
    };
  }

  @Override
  public int initEncrypt(ByteBuffer output) throws KeyczarException {
    try {
      cipher.init(Cipher.ENCRYPT_MODE, jcePublicKey);
    } catch (InvalidKeyException e) {
      throw new KeyczarException(e);
    }
    return 0;
  }

  @Override
  public int maxOutputSize(int inputLen) {
    return digestSize();
  }

  @Override
  public int updateEncrypt(ByteBuffer input, ByteBuffer output)
      throws KeyczarException {
    try {
      return cipher.update(input, output);
    } catch (ShortBufferException e) {
      throw new KeyczarException(e);
    }
  }

}
